/-
Copyright (c) 2019 Zhouhang Zhou. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Zhouhang Zhou, Frédéric Dupuis, Heather Macbeth
-/
module

public import Mathlib.Analysis.InnerProductSpace.Projection.Basic
public import Mathlib.Analysis.InnerProductSpace.Projection.Reflection
public import Mathlib.Analysis.InnerProductSpace.Projection.Submodule
public import Mathlib.Algebra.DirectSum.Decomposition
public import Mathlib.Analysis.Normed.Module.RCLike.Basic

/-!
# Orthogonal projections in finite-dimensional spaces

This file contains results about orthogonal projections in finite-dimensional spaces.

## Main results

* `Submodule.det_reflection`: The determinant of `K.reflection` is `(-1) ^ finrank 𝕜 Kᗮ`.
* `LinearIsometryEquiv.reflections_generate`: The orthogonal group of `F` is
generated by reflections.

## Results that do not use finite dimensionality:

* `OrthogonalFamily.isInternal_iff_of_isComplete`
* `OrthogonalFamily.sum_projection_of_mem_iSup`
* `OrthogonalFamily.projection_directSum_coeAddHom`
* `OrthogonalFamily.decomposition`
* `maximal_orthonormal_iff_orthogonalComplement_eq_bot`
-/

@[expose] public section

variable {𝕜 E F : Type*} [RCLike 𝕜]
variable [NormedAddCommGroup E] [NormedAddCommGroup F]
variable [InnerProductSpace 𝕜 E] [InnerProductSpace ℝ F]

local notation "⟪" x ", " y "⟫" => inner 𝕜 x y
local notation "absR" => @abs ℝ _ _

open Topology Finsupp Submodule RCLike Real Filter InnerProductSpace
open LinearMap (ker range)

variable (K : Submodule 𝕜 E)

namespace Submodule

section FiniteDimensional

open Module

variable [FiniteDimensional 𝕜 K]

@[simp]
theorem topologicalClosure_eq_self : K.topologicalClosure = K :=
  K.closed_of_finiteDimensional.submodule_topologicalClosure_eq

@[simp]
theorem det_reflection : LinearMap.det K.reflection.toLinearMap = (-1) ^ finrank 𝕜 Kᗮ := by
  by_cases hK : FiniteDimensional 𝕜 Kᗮ
  swap
  · rw [finrank_of_infinite_dimensional hK, pow_zero, LinearMap.det_eq_one_of_finrank_eq_zero]
    exact finrank_of_infinite_dimensional fun h ↦ hK (h.finiteDimensional_submodule _)
  let e := K.prodEquivOfIsCompl _ K.isCompl_orthogonal_of_hasOrthogonalProjection
  let b := (finBasis 𝕜 K).prod (finBasis 𝕜 Kᗮ)
  have : LinearMap.toMatrix b b (e.symm ∘ₗ K.reflection.toLinearMap ∘ₗ e.symm.symm) =
      Matrix.fromBlocks 1 0 0 (-1) := by
    ext (_ | _) (_ | _) <;>
    simp [LinearMap.toMatrix_apply, b, Matrix.one_apply, Finsupp.single_apply, e, eq_comm,
      reflection_mem_subspace_eq_self, reflection_mem_subspace_orthogonalComplement_eq_neg]
  rw [← LinearMap.det_conj _ e.symm, ← LinearMap.det_toMatrix b, this, Matrix.det_fromBlocks_zero₂₁,
    Matrix.det_one, one_mul, Matrix.det_neg, Fintype.card_fin, Matrix.det_one, mul_one]

@[simp]
theorem linearEquiv_det_reflection : K.reflection.det = (-1) ^ finrank 𝕜 Kᗮ := by
  ext
  rw [LinearEquiv.coe_det, Units.val_pow_eq_pow_val]
  exact K.det_reflection

end FiniteDimensional

variable {K}

open Module

/-- Given a finite-dimensional subspace `K₂`, and a subspace `K₁`
contained in it, the dimensions of `K₁` and the intersection of its
orthogonal subspace with `K₂` add to that of `K₂`. -/
theorem finrank_add_inf_finrank_orthogonal {K₁ K₂ : Submodule 𝕜 E}
    [FiniteDimensional 𝕜 K₂] (h : K₁ ≤ K₂) :
    finrank 𝕜 K₁ + finrank 𝕜 (K₁ᗮ ⊓ K₂ : Submodule 𝕜 E) = finrank 𝕜 K₂ := by
  haveI : FiniteDimensional 𝕜 K₁ := Submodule.finiteDimensional_of_le h
  haveI := FiniteDimensional.proper_rclike 𝕜 K₁
  have hd := Submodule.finrank_sup_add_finrank_inf_eq K₁ (K₁ᗮ ⊓ K₂)
  rw [← inf_assoc, (Submodule.orthogonal_disjoint K₁).eq_bot, bot_inf_eq, finrank_bot,
    Submodule.sup_orthogonal_inf_of_hasOrthogonalProjection h] at hd
  rw [add_zero] at hd
  exact hd.symm

/-- Given a finite-dimensional subspace `K₂`, and a subspace `K₁`
contained in it, the dimensions of `K₁` and the intersection of its
orthogonal subspace with `K₂` add to that of `K₂`. -/
theorem finrank_add_inf_finrank_orthogonal' {K₁ K₂ : Submodule 𝕜 E}
    [FiniteDimensional 𝕜 K₂] (h : K₁ ≤ K₂) {n : ℕ} (h_dim : finrank 𝕜 K₁ + n = finrank 𝕜 K₂) :
    finrank 𝕜 (K₁ᗮ ⊓ K₂ : Submodule 𝕜 E) = n := by
  rw [← add_right_inj (finrank 𝕜 K₁)]
  simp [Submodule.finrank_add_inf_finrank_orthogonal h, h_dim]

/-- Given a finite-dimensional space `E` and subspace `K`, the dimensions of `K` and `Kᗮ` add to
that of `E`. -/
theorem finrank_add_finrank_orthogonal [FiniteDimensional 𝕜 E] (K : Submodule 𝕜 E) :
    finrank 𝕜 K + finrank 𝕜 Kᗮ = finrank 𝕜 E := by
  convert Submodule.finrank_add_inf_finrank_orthogonal (le_top : K ≤ ⊤) using 1
  · rw [inf_top_eq]
  · simp

/-- Given a finite-dimensional space `E` and subspace `K`, the dimensions of `K` and `Kᗮ` add to
that of `E`. -/
theorem finrank_add_finrank_orthogonal' [FiniteDimensional 𝕜 E] {K : Submodule 𝕜 E}
    {n : ℕ} (h_dim : finrank 𝕜 K + n = finrank 𝕜 E) : finrank 𝕜 Kᗮ = n := by
  rw [← add_right_inj (finrank 𝕜 K)]
  simp [Submodule.finrank_add_finrank_orthogonal, h_dim]

/-- In a finite-dimensional inner product space, the dimension of the orthogonal complement of the
span of a nonzero vector is one less than the dimension of the space. -/
theorem finrank_orthogonal_span_singleton {n : ℕ} [_i : Fact (finrank 𝕜 E = n + 1)] {v : E}
    (hv : v ≠ 0) : finrank 𝕜 (𝕜 ∙ v)ᗮ = n := by
  haveI : FiniteDimensional 𝕜 E := .of_fact_finrank_eq_succ n
  exact finrank_add_finrank_orthogonal' <| by
    simp [finrank_span_singleton hv, _i.elim, add_comm]

end Submodule

open Module Submodule

/-- An element `φ` of the orthogonal group of `F` can be factored as a product of reflections, and
specifically at most as many reflections as the dimension of the complement of the fixed subspace
of `φ`. -/
theorem LinearIsometryEquiv.reflections_generate_dim_aux [FiniteDimensional ℝ F] {n : ℕ}
    (φ : F ≃ₗᵢ[ℝ] F) (hn : finrank ℝ (ker (ContinuousLinearMap.id ℝ F - φ))ᗮ ≤ n) :
    ∃ l : List F, l.length ≤ n ∧ φ = (l.map fun v => (ℝ ∙ v)ᗮ.reflection).prod := by
  -- We prove this by strong induction on `n`, the dimension of the orthogonal complement of the
  -- fixed subspace of the endomorphism `φ`
  induction n generalizing φ with
  | zero => -- Base case: `n = 0`, the fixed subspace is the whole space, so `φ = id`
    refine ⟨[], rfl.le, show φ = 1 from ?_⟩
    have : ker (ContinuousLinearMap.id ℝ F - φ) = ⊤ := by
      rwa [le_zero_iff, finrank_eq_zero, orthogonal_eq_bot_iff] at hn
    symm
    ext x
    have := LinearMap.congr_fun (LinearMap.ker_eq_top.mp this) x
    simpa only [sub_eq_zero, ContinuousLinearMap.coe_sub, LinearMap.sub_apply,
      LinearMap.zero_apply] using this
  | succ n IH =>
    -- Inductive step.  Let `W` be the fixed subspace of `φ`.  We suppose its complement to have
    -- dimension at most n + 1.
    let W := ker (ContinuousLinearMap.id ℝ F - φ)
    have hW : ∀ w ∈ W, φ w = w := fun w hw => (sub_eq_zero.mp hw).symm
    by_cases hn' : finrank ℝ Wᗮ ≤ n
    · obtain ⟨V, hV₁, hV₂⟩ := IH φ hn'
      exact ⟨V, hV₁.trans n.le_succ, hV₂⟩
    -- Take a nonzero element `v` of the orthogonal complement of `W`.
    haveI : Nontrivial Wᗮ := nontrivial_of_finrank_pos (by lia : 0 < finrank ℝ Wᗮ)
    obtain ⟨v, hv⟩ := exists_ne (0 : Wᗮ)
    have hφv : φ v ∈ Wᗮ := by
      intro w hw
      rw [← hW w hw, LinearIsometryEquiv.inner_map_map]
      exact v.prop w hw
    have hv' : (v : F) ∉ W := by
      intro h
      exact hv ((mem_left_iff_eq_zero_of_disjoint W.orthogonal_disjoint).mp h)
    -- Let `ρ` be the reflection in `v - φ v`; this is designed to swap `v` and `φ v`
    let x : F := v - φ v
    let ρ := (ℝ ∙ x)ᗮ.reflection
    -- Notation: Let `V` be the fixed subspace of `φ.trans ρ`
    let V := ker (ContinuousLinearMap.id ℝ F - φ.trans ρ)
    have hV : ∀ w, ρ (φ w) = w → w ∈ V := by
      intro w hw
      change w - ρ (φ w) = 0
      rw [sub_eq_zero, hw]
    -- Everything fixed by `φ` is fixed by `φ.trans ρ`
    have H₂V : W ≤ V := by
      intro w hw
      apply hV
      rw [hW w hw]
      refine reflection_mem_subspace_eq_self ?_
      rw [mem_orthogonal_singleton_iff_inner_left]
      exact Submodule.sub_mem _ v.prop hφv _ hw
    -- `v` is also fixed by `φ.trans ρ`
    have H₁V : (v : F) ∈ V := by
      apply hV
      have : ρ v = φ v := reflection_sub (φ.norm_map v).symm
      rw [← this]
      exact reflection_reflection _ _
    -- By dimension-counting, the complement of the fixed subspace of `φ.trans ρ` has dimension at
    -- most `n`
    have : finrank ℝ Vᗮ ≤ n := by
      change finrank ℝ Wᗮ ≤ n + 1 at hn
      have : finrank ℝ W + 1 ≤ finrank ℝ V :=
        finrank_lt_finrank_of_lt (SetLike.lt_iff_le_and_exists.2 ⟨H₂V, v, H₁V, hv'⟩)
      have : finrank ℝ V + finrank ℝ Vᗮ = finrank ℝ F := V.finrank_add_finrank_orthogonal
      have : finrank ℝ W + finrank ℝ Wᗮ = finrank ℝ F := W.finrank_add_finrank_orthogonal
      omega
    -- So apply the inductive hypothesis to `φ.trans ρ`
    obtain ⟨l, hl, hφl⟩ := IH (ρ * φ) this
    -- Prepend `ρ` to the factorization into reflections obtained for `φ.trans ρ`; this gives a
    -- factorization into reflections for `φ`.
    refine ⟨x::l, Nat.succ_le_succ hl, ?_⟩
    rw [List.map_cons, List.prod_cons]
    have := congr_arg (ρ * ·) hφl
    dsimp only at this
    rwa [← mul_assoc, reflection_mul_reflection, one_mul] at this

/-- The orthogonal group of `F` is generated by reflections; specifically each element `φ` of the
orthogonal group is a product of at most as many reflections as the dimension of `F`.

Special case of the **Cartan–Dieudonné theorem**. -/
theorem LinearIsometryEquiv.reflections_generate_dim [FiniteDimensional ℝ F] (φ : F ≃ₗᵢ[ℝ] F) :
    ∃ l : List F, l.length ≤ finrank ℝ F ∧ φ = (l.map fun v => reflection (ℝ ∙ v)ᗮ).prod :=
  let ⟨l, hl₁, hl₂⟩ := φ.reflections_generate_dim_aux le_rfl
  ⟨l, hl₁.trans (finrank_le _), hl₂⟩

/-- The orthogonal group of `F` is generated by reflections. -/
theorem LinearIsometryEquiv.reflections_generate [FiniteDimensional ℝ F] :
    Subgroup.closure (Set.range fun v : F => reflection (ℝ ∙ v)ᗮ) = ⊤ := by
  rw [Subgroup.eq_top_iff']
  intro φ
  rcases φ.reflections_generate_dim with ⟨l, _, rfl⟩
  apply (Subgroup.closure _).list_prod_mem
  intro x hx
  rcases List.mem_map.mp hx with ⟨a, _, hax⟩
  exact Subgroup.subset_closure ⟨a, hax⟩

section OrthogonalFamily

open Submodule

variable {ι : Type*}

/-- An orthogonal family of subspaces of `E` satisfies `DirectSum.IsInternal` (that is,
they provide an internal direct sum decomposition of `E`) if and only if their span has trivial
orthogonal complement. -/
theorem OrthogonalFamily.isInternal_iff_of_isComplete [DecidableEq ι] {V : ι → Submodule 𝕜 E}
    (hV : OrthogonalFamily 𝕜 (fun i => V i) fun i => (V i).subtypeₗᵢ)
    (hc : IsComplete (↑(iSup V) : Set E)) : DirectSum.IsInternal V ↔ (iSup V)ᗮ = ⊥ := by
  haveI : CompleteSpace (↥(iSup V)) := hc.completeSpace_coe
  simp only [DirectSum.isInternal_submodule_iff_iSupIndep_and_iSup_eq_top, hV.independent,
    true_and, orthogonal_eq_bot_iff]

/-- An orthogonal family of subspaces of `E` satisfies `DirectSum.IsInternal` (that is,
they provide an internal direct sum decomposition of `E`) if and only if their span has trivial
orthogonal complement. -/
theorem OrthogonalFamily.isInternal_iff [DecidableEq ι] [FiniteDimensional 𝕜 E]
    {V : ι → Submodule 𝕜 E} (hV : OrthogonalFamily 𝕜 (fun i => V i) fun i => (V i).subtypeₗᵢ) :
    DirectSum.IsInternal V ↔ (iSup V)ᗮ = ⊥ :=
  haveI := FiniteDimensional.proper_rclike 𝕜 (↥(iSup V))
  hV.isInternal_iff_of_isComplete (completeSpace_coe_iff_isComplete.mp inferInstance)

open DirectSum

/-- If `x` lies within an orthogonal family `v`, it can be expressed as a sum of projections. -/
theorem OrthogonalFamily.sum_projection_of_mem_iSup [Fintype ι] {V : ι → Submodule 𝕜 E}
    [∀ i, CompleteSpace (V i)] (hV : OrthogonalFamily 𝕜 (fun i => V i) fun i => (V i).subtypeₗᵢ)
    (x : E) (hx : x ∈ iSup V) : (∑ i, (V i).starProjection x) = x := by
  induction hx using iSup_induction' with
  | mem i x hx =>
    refine
      (Finset.sum_eq_single_of_mem i (Finset.mem_univ _) fun j _ hij => ?_).trans
        (starProjection_eq_self_iff.mpr hx)
    rw [starProjection_apply, orthogonalProjection_mem_subspace_orthogonalComplement_eq_zero,
      Submodule.coe_zero]
    exact hV.isOrtho hij.symm hx
  | zero =>
    simp_rw [map_zero, Finset.sum_const_zero]
  | add x y _ _ hx hy =>
    simp_rw [map_add, Finset.sum_add_distrib]
    exact congr_arg₂ (· + ·) hx hy

/-- If a family of submodules is orthogonal, then the `orthogonalProjection` on a direct sum
is just the coefficient of that direct sum. -/
theorem OrthogonalFamily.projection_directSum_coeAddHom [DecidableEq ι] {V : ι → Submodule 𝕜 E}
    (hV : OrthogonalFamily 𝕜 (fun i => V i) fun i => (V i).subtypeₗᵢ) (x : ⨁ i, V i) (i : ι)
    [CompleteSpace (V i)] :
    (V i).orthogonalProjection (DirectSum.coeAddMonoidHom V x) = x i := by
  induction x using DirectSum.induction_on with
  | zero => simp
  | of j x =>
    simp_rw [DirectSum.coeAddMonoidHom_of, DirectSum.of,
      -- Need to unfold `DirectSum` to see through the defeq abuse.
      DirectSum, DFinsupp.singleAddHom_apply]
    obtain rfl | hij := Decidable.eq_or_ne i j
    · rw [orthogonalProjection_mem_subspace_eq_self, DFinsupp.single_eq_same]
    · rw [orthogonalProjection_mem_subspace_orthogonalComplement_eq_zero,
        DFinsupp.single_eq_of_ne hij]
      exact hV.isOrtho hij.symm x.prop
  | add x y hx hy =>
    simp_rw [map_add]
    exact congr_arg₂ (· + ·) hx hy

/-- If a family of submodules is orthogonal and they span the whole space, then the orthogonal
projection provides a means to decompose the space into its submodules.

The projection function is `decompose V x i = (V i).orthogonalProjection x`.

See note [reducible non-instances]. -/
noncomputable abbrev OrthogonalFamily.decomposition
    [DecidableEq ι] [Fintype ι] {V : ι → Submodule 𝕜 E}
    [∀ i, CompleteSpace (V i)] (hV : OrthogonalFamily 𝕜 (fun i => V i) fun i => (V i).subtypeₗᵢ)
    (h : iSup V = ⊤) : DirectSum.Decomposition V where
  decompose' x := DFinsupp.equivFunOnFintype.symm fun i => (V i).orthogonalProjection x
  left_inv x := by
    dsimp only
    letI := fun i => Classical.decEq (V i)
    rw [DirectSum.coeAddMonoidHom, DirectSum.toAddMonoid, DFinsupp.liftAddHom_apply]
    -- This used to be `rw`, but we need `erw` after https://github.com/leanprover/lean4/pull/2644
    erw [DFinsupp.sumAddHom_apply]; rw [DFinsupp.sum_eq_sum_fintype]
    · simp_rw [Equiv.apply_symm_apply, AddSubmonoidClass.coe_subtype]
      exact hV.sum_projection_of_mem_iSup _ ((h.ge :) Submodule.mem_top)
    · intro i
      exact map_zero _
  right_inv x := by
    dsimp only
    simp_rw [hV.projection_directSum_coeAddHom, DFinsupp.equivFunOnFintype_symm_coe]

end OrthogonalFamily

section OrthonormalBasis

variable {v : Set E}

open Module Submodule Set

/-- An orthonormal set in an `InnerProductSpace` is maximal, if and only if the orthogonal
complement of its span is empty. -/
theorem maximal_orthonormal_iff_orthogonalComplement_eq_bot (hv : Orthonormal 𝕜 ((↑) : v → E)) :
    (∀ u ⊇ v, Orthonormal 𝕜 ((↑) : u → E) → u = v) ↔ (span 𝕜 v)ᗮ = ⊥ := by
  rw [Submodule.eq_bot_iff]
  constructor
  · contrapose!
    -- ** direction 1: nonempty orthogonal complement implies nonmaximal
    rintro ⟨x, hx', hx⟩
    -- take a nonzero vector and normalize it
    let e := (‖x‖⁻¹ : 𝕜) • x
    have he : ‖e‖ = 1 := by simp [e, norm_smul_inv_norm hx]
    have he' : e ∈ (span 𝕜 v)ᗮ := smul_mem' _ _ hx'
    have he'' : e ∉ v := by
      intro hev
      have : e = 0 := by
        have : e ∈ span 𝕜 v ⊓ (span 𝕜 v)ᗮ := ⟨subset_span hev, he'⟩
        simpa [(span 𝕜 v).inf_orthogonal_eq_bot] using this
      have : e ≠ 0 := hv.ne_zero ⟨e, hev⟩
      contradiction
    -- put this together with `v` to provide a candidate orthonormal basis for the whole space
    refine ⟨insert e v, v.subset_insert e, ⟨?_, ?_⟩, (ne_insert_of_notMem v he'').symm⟩
    · -- show that the elements of `insert e v` have unit length
      rintro ⟨a, ha'⟩
      rcases eq_or_mem_of_mem_insert ha' with ha | ha
      · simp [ha, he]
      · exact hv.1 ⟨a, ha⟩
    · -- show that the elements of `insert e v` are orthogonal
      have h_end : ∀ a ∈ v, ⟪a, e⟫ = 0 := by
        intro a ha
        exact he' a (Submodule.subset_span ha)
      rintro ⟨a, ha'⟩
      rcases eq_or_mem_of_mem_insert ha' with ha | ha
      · rintro ⟨b, hb'⟩ hab'
        have hb : b ∈ v := by grind
        rw [inner_eq_zero_symm]
        simpa [ha] using h_end b hb
      rintro ⟨b, hb'⟩ hab'
      rcases eq_or_mem_of_mem_insert hb' with hb | hb
      · simpa [hb] using h_end a ha
      have : (⟨a, ha⟩ : v) ≠ ⟨b, hb⟩ := by
        intro hab''
        apply hab'
        simpa using hab''
      exact hv.2 this
  · -- ** direction 2: empty orthogonal complement implies maximal
    simp only [Subset.antisymm_iff]
    rintro h u (huv : v ⊆ u) hu
    refine ⟨?_, huv⟩
    intro x hxu
    refine ((mt (h x)) (hu.ne_zero ⟨x, hxu⟩)).imp_symm ?_
    intro hxv y hy
    have hxv' : (⟨x, hxu⟩ : u) ∉ ((↑) ⁻¹' v : Set u) := by simp [hxv]
    obtain ⟨l, hl, rfl⟩ :
      ∃ l ∈ supported 𝕜 𝕜 ((↑) ⁻¹' v : Set u), (linearCombination 𝕜 ((↑) : u → E)) l = y := by
      rw [← Finsupp.mem_span_image_iff_linearCombination]
      simp [huv, inter_eq_self_of_subset_right, hy]
    exact hu.inner_finsupp_eq_zero hxv' hl

variable [FiniteDimensional 𝕜 E]

/-- An orthonormal set in a finite-dimensional `InnerProductSpace` is maximal, if and only if it
is a basis. -/
theorem maximal_orthonormal_iff_basis_of_finiteDimensional (hv : Orthonormal 𝕜 ((↑) : v → E)) :
    (∀ u ⊇ v, Orthonormal 𝕜 ((↑) : u → E) → u = v) ↔ ∃ b : Basis v 𝕜 E, ⇑b = ((↑) : v → E) := by
  haveI := FiniteDimensional.proper_rclike 𝕜 (span 𝕜 v)
  rw [maximal_orthonormal_iff_orthogonalComplement_eq_bot hv]
  rw [Submodule.orthogonal_eq_bot_iff]
  have hv_coe : range ((↑) : v → E) = v := by simp
  constructor
  · refine fun h => ⟨Basis.mk hv.linearIndependent _, Basis.coe_mk _ ?_⟩
    convert h.ge
  · rintro ⟨h, coe_h⟩
    rw [← h.span_eq, coe_h, hv_coe]

end OrthonormalBasis
